"""识别算法实现"""

import numpy as np
import xarray as xr
from xarray import DataArray
from typing import List, Dict
from scipy.signal import argrelextrema
from scipy.optimize import minimize
import pycwt as wavelet
from config import IdentifyConfig as cfg
from mrgLike import MrgwLikeEvent, MrgwLikeEvents, MrgwLike
import pandas as pd

def mrg_err_signle_lon(y0: float, v: np.ndarray, eqV: float,
                        lat: np.ndarray) -> float:
    """单个经度上的MRG误差，用于optimize.minimize的目标函数"""
    vFitted = eqV * np.exp(-(lat/y0)**2)
    return np.sqrt(np.mean((v - vFitted)**2)) / np.abs(eqV)


def identify_mrgwl_event(flt_vdata: DataArray, lon_name: str = "longitude",
                 lat_name: str = "latitude", time_name: str="time",
                 ) -> MrgwLikeEvents:
    """识别MRGWL事件，需要传入滤波后的经向风数据，时间长度不限。

    Parameters
    ----------
    flt_vdata : DataArray
        滤波后的850 hPa 经向风数据，需要南北半球都有数据。
    lon_name : str, optional
        数据中经度为名称, by default "longitude"
    lat_name : str, optional
        数据中纬度的名称, by default "latitude"
    time_name : str, optional
        数据中时间的名称, by default "time"

    Returns
    -------
    MrgwLikeEvents
        识别到的MRG事件列表对象
    """

    data_res = abs(flt_vdata[lon_name].values[1] - flt_vdata[lon_name].values[0])

    concat_number = int(cfg.data_lat_range / data_res)

    ext_order = int(cfg.min_mrg_wave_width / 2 / data_res)

    # 选择范围
    flt_vdata = flt_vdata.sel({lat_name: slice(cfg.data_lat_range,
                                            -cfg.data_lat_range)})

    # 1. 计算对称部分
    data_north = flt_vdata.sel({lat_name: slice(None, 0)})
    data_south = flt_vdata.sel({lat_name: slice(0, None)})

    # 反转南半球的纬度
    data_south = data_south.assign_coords({lat_name: -data_south.latitude.values})
    data = (data_north + data_south) / 2
    latitude = data[lat_name].values
    longitude = data[lon_name].values

    # 2. 找出赤道上经向风的局地极大极小值

    mrg_results = MrgwLikeEvents()

    # 开始时间历遍，这里效率会比较低，有待优化
    for t_idx in range(data[time_name].size):
        mrg_result_cache: List[MrgwLikeEvent] = []
        time = data[time_name].values[t_idx]
        vwnd = data.isel({time_name: t_idx})
        eqV = vwnd.sel({lat_name: 0})

        local_max = argrelextrema(eqV.values, np.greater, 
                                    order=ext_order, axis=-1)[0]
        local_min = argrelextrema(eqV.values, np.less,
                                    order=ext_order, axis=-1)[0]
        local_max_min = np.concatenate([local_max, local_min], axis=-1)
        local_max_min.sort()

        # 3. 计算每一处极值的最优 y0 和 误差
        for idx in local_max_min:
            res = minimize(mrg_err_signle_lon, 6,
                           args=(vwnd.values[:, idx], eqV.values[idx], latitude),
                           bounds=[(1e-1, cfg.y0_threshold)])
            if res.fun < cfg.err_threshold_first:
                mrg_result_cache.append(
                    MrgwLikeEvent(
                        eq_v=eqV.values[idx], y0=res.x[0], err=res.fun,
                        time=time, longitude=longitude[idx], wavenumber=0,
                        lon_idx=idx
                    )
                )

        # 4. 对赤道风小波分析，经度循环拼接
        if len(mrg_result_cache) > 0:
            eq_v_concat = np.concatenate([eqV.values[-concat_number:], eqV.values,
                                          eqV.values[concat_number:]])
            freqs = np.arange(cfg.low_pass_wavenumber, 1-cfg.cwt_wavenumber_precision,
                            -cfg.cwt_wavenumber_precision)
            wave, _, freqs, _, _, _ = wavelet.cwt(eq_v_concat,
                        data_res, wavelet='mexicanhat', freqs=1/(360/freqs))
            wave = wave[:, concat_number:-concat_number]

            for mrg in mrg_result_cache:
                lon_freq = np.argmax(np.abs(wave[:, mrg._lon_idx]))
                lon_freq = freqs[lon_freq]
                mrg.set_wavenumber(round(lon_freq*360, 2))

                mrg_results.append_event(mrg)

    # 到此得到了这个数据中每一个时刻的识别到的MRG波的属性
    return mrg_results


def lon_time_distance(lon1: float, lon2: float, 
                      time1: pd.Timestamp, time2: pd.Timestamp,
                      timeCoef=1) -> float:
    """计算两个时次的MRG事件的距离，如果大于最大速度或者在东边返回999，
    考虑时间差，以小时为单位，乘以timeCoef作为y向距离，然后算欧拉距离"""

    def distance(l1, l2, t1, t2, timeCoef, dt):
        return np.sqrt((l1 - l2)**2 + (dt*timeCoef)**2)
    
    dt = (time2 - time1).total_seconds() / 3600
    max_distance = dt * cfg.max_speed

    # 如果正好在360度的两边
    if lon1 <= max_distance and lon2 >= 360 - max_distance:
        dis = distance(lon1+360, lon2, time1, time2, timeCoef, dt)
        if dis <= max_distance:
            return dis
        else:
            return 999
    # 东传
    elif lon2 > lon1:
        return 999
    # 西传
    else:
        dis = distance(lon1, lon2, time1, time2, timeCoef, dt)
        if dis <= max_distance:
            return dis
        else:
            return 999


def get_two_line_closest(line0: pd.DataFrame, line1: pd.DataFrame, 
                         distance=False) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
    """获取两个时间中 MRGWL 事件两两的距离并返回从近到远的索引对。
    """
    distanceList = np.full((len(line0), len(line1)), 999, dtype='float32')
    idx0 = line0.index
    idx1 = line1.index

    # 计算距离
    for i, a in enumerate(idx0):
        for j, b in enumerate(idx1):
            # 如果相位不同
            if np.sign(line0.loc[a]['eq_v']) != np.sign(line1.loc[b]['eq_v']):
                continue

            distanceList[i, j] = lon_time_distance(
                line0.loc[a]['longitude'],
                line1.loc[b]['longitude'],
                line0.loc[a]['time'],
                line1.loc[b]['time'],
            )

    # 从距离最近的开始选
    min_index = np.unravel_index(
        np.argsort(distanceList, axis=None), distanceList.shape)
    
    currentLineSelected = np.zeros(len(line0), dtype='bool')
    followingLineSelected = np.zeros(len(idx1), dtype='bool')

    indexResult = []
    if distance:
        distance_result = []
    for i in range(len(min_index[0])):
        dis = distanceList[min_index[0][i], min_index[1][i]]
        if dis >= 999:
            break

        # 如果已经被选则跳过
        if currentLineSelected[min_index[0][i]] or \
            followingLineSelected[min_index[1][i]]:
            continue

        indexResult.append((idx0[min_index[0][i]], idx1[min_index[1][i]]))
        if distance:
            distance_result.append(dis)

        currentLineSelected[min_index[0][i]] = True
        followingLineSelected[min_index[1][i]] = True

    if distance:
        return indexResult, distance_result
    else:
        return indexResult


def mrgwl_case_identification(mrgwl_events: MrgwLikeEvents,
                              error_threshold: float = 0.1,
                              min_exist_hours: float = cfg.min_exist_hours) -> list[MrgwLike]:
    """将识别到的 MRGWL 事件进行连接，形成连续的 MRGWL 传播过程

    Parameters
    ----------
    mrgwl_events : MrgwLikeEvents
        识别到的MRGWL事件列表
    error_threshold : float, optional
        误差限制，误差大于此的 MRGWL 事件将被忽略, by default 0.1
    min_exist_hours : float, optional
        最短存在时间, by default cfg.min_exist_hours

    Returns
    -------
    list[MrgwLike]
        连接后的 MRGWL 列表
    """
    mrg_df = mrgwl_events.df
    mrg_df = mrg_df[mrg_df['err'] < error_threshold]

    # 1. 按照时间分组
    mrgwl_group = mrg_df.groupby('time')
    time_list = list(mrgwl_group.groups.keys())
    index_list = list(mrgwl_group.groups.values())

    mrgwl_id = 0
    mrgwl_id_list = {i: -1 for i in mrg_df.index}

    mrgwl_result: Dict[int, MrgwLike] = {}

    # 2. 连接
    for i in range(len(time_list) - 2):
        time = time_list[i]
        line0 = mrg_df.loc[index_list[i]]
        line1 = mrg_df.loc[index_list[i+1]]
        line2 = mrg_df.loc[index_list[i+2]]

        # 寻找下一行
        dis1 = get_two_line_closest(line0, line1)
        checked_idx = []

        for j in range(len(dis1)):
            idx0 = dis1[j][0]
            idx1 = dis1[j][1]
            checked_idx.append(idx0)
            if mrgwl_id_list[idx0] < 0:
                # 没有记录，创建
                mrgwl_id_list[idx0] = mrgwl_id
                mrgwl_result[mrgwl_id] = MrgwLike(id=mrgwl_id, err_limit=error_threshold)
                mrgwl_result[mrgwl_id].append_event(line0.loc[idx0])
                mrgwl_id += 1
            
            # 添加下一行的
            mid = mrgwl_id_list[idx0]
            mrgwl_result[mid].append_event(line1.loc[idx1])
            mrgwl_id_list[idx1] = mid
        
        # 检查是否有未被选中的
        not_checked_mrgwl_idx = np.array([i for i in line0.index
                                            if not i in checked_idx])

        # 如果没有未被选中的,跳过。还有没选中的，看看能不能与下下行连接
        if len(not_checked_mrgwl_idx) == 0:
            continue
        
        line0_not_checked = mrg_df.loc[not_checked_mrgwl_idx]

        # 计算到下下行的距离
        dis_idx2, dis2 = get_two_line_closest(line0_not_checked, line2, distance=True)

        # 如果没有距离下下行近的，跳过
        if len(dis2) == 0:
            continue
        
        # 计算下行到下下行的距离
        dis_idx3, dis3 = get_two_line_closest(line1, line2, distance=True)

        # 检查哪一行距离更近，优先分配给近的
        for j in range(len(dis2)):
            idx0 = dis_idx2[j][0]
            idx2 = dis_idx2[j][1]
            dis02 = dis2[j]
            dis12 = 999
            for k in range(len(dis3)):
                # 也距离
                if dis_idx3[k][1] == idx2:
                    dis12 = dis3[k]
                    break
            if dis02 < dis12:
                # 如果当前行的没有被识别为 MGRWL
                if mrgwl_id_list[idx0] < 0:
                    mrgwl_result[mrgwl_id] = MrgwLike(id=mrgwl_id, err_limit=error_threshold)
                    mrgwl_result[mrgwl_id].append_event(line0.loc[idx0])
                    mrgwl_id_list[idx0] = mrgwl_id
                    mrgwl_id += 1

                # 添加下下行的
                mid = mrgwl_id_list[idx0]
                mrgwl_result[mid].append_event(line2.loc[idx2])
                mrgwl_id_list[idx2] = mid

    # 3. 过滤、计算相速度
    valid_mrgwl = []
    mrgwl_id = 0
    for i in range(len(mrgwl_result)):
        mrgwl = mrgwl_result[i]
        if len(mrgwl.df) < 2:
            continue
        if (mrgwl.df['time'].values[-1] - mrgwl.df['time'].values[0]).total_seconds() / 3600 < min_exist_hours:
            continue

        mrgwl.cal_phase_speed()
        mrgwl.set_uid()

        mrgwl.id = mrgwl_id
        mrgwl_id += 1
        valid_mrgwl.append(mrgwl)

    return valid_mrgwl